from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("GnuTLS MCP Server")

"""
GnuTLS MCP Server

面向软件：gnutls (GNU Transport Layer Security Library)

MCP tools 列表：
1. gnutls_check_installed - 检查 gnutls 是否已安装
   参数：无
   返回：包含安装状态和执行结果的 JSON

2. gnutls_connect - 建立到远程服务器的 TLS 连接
   参数：host (str) - 目标主机名或IP地址
   返回：包含连接结果的 JSON

3. gnutls_get_version - 获取 gnutls-cli 版本信息
   参数：无
   返回：包含版本信息的 JSON

4. gnutls_get_certificate - 获取远程服务器的证书信息
   参数：host (str) - 目标主机名或IP地址
   返回：包含证书信息的 JSON

5. gnutls_handshake - 执行 TLS 握手并验证完成状态
   参数：host (str) - 目标主机名或IP地址
   返回：包含握手结果的 JSON

6. gnutls_list_ciphers - 列出可用的密码套件
   参数：无
   返回：包含密码套件列表的 JSON

7. gnutls_list_priorities - 显示优先级列表
   参数：无
   返回：包含优先级列表的 JSON

8. gnutls_benchmark_ciphers - 对密码套件进行基准测试
   参数：无
   返回：包含基准测试结果的 JSON

9. gnutls_save_certificate - 保存远程服务器的证书到文件
   参数：host (str) - 目标主机名或IP地址, cert_file (str) - 证书保存路径
   返回：包含保存结果的 JSON
"""

@mcp.tool()
def gnutls_check_installed(ctx: Context) -> dict:
    """检查 gnutls 是否已安装
    
    参数：无
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["rpm", "-qa", "|", "grep", "gnutls"]
    result = subprocess.run(" ".join(cmd), shell=True, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_connect(ctx: Context, host: str) -> dict:
    """建立到远程服务器的 TLS 连接
    
    参数：
    host (str): 目标主机名或IP地址
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--insecure", host]
    result = subprocess.run(cmd, capture_output=True, text=True, input="", timeout=30)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_get_version(ctx: Context) -> dict:
    """获取 gnutls-cli 版本信息
    
    参数：无
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--version"]
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_get_certificate(ctx: Context, host: str) -> dict:
    """获取远程服务器的证书信息
    
    参数：
    host (str): 目标主机名或IP地址
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--print-cert", "--insecure", host]
    result = subprocess.run(cmd, capture_output=True, text=True, input="", timeout=30)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_handshake(ctx: Context, host: str) -> dict:
    """执行 TLS 握手并验证完成状态
    
    参数：
    host (str): 目标主机名或IP地址
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--priority", "NORMAL", "--insecure", host]
    result = subprocess.run(cmd, capture_output=True, text=True, input="", timeout=30)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_list_ciphers(ctx: Context) -> dict:
    """列出可用的密码套件
    
    参数：无
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--list"]
    result = subprocess.run(cmd, capture_output=True, text=True, input="")
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_list_priorities(ctx: Context) -> dict:
    """显示优先级列表
    
    参数：无
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--priority-list"]
    result = subprocess.run(cmd, capture_output=True, text=True, input="")
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_benchmark_ciphers(ctx: Context) -> dict:
    """对密码套件进行基准测试
    
    参数：无
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--benchmark-ciphers"]
    result = subprocess.run(cmd, capture_output=True, text=True, input="")
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def gnutls_save_certificate(ctx: Context, host: str, cert_file: str) -> dict:
    """保存远程服务器的证书到文件
    
    参数：
    host (str): 目标主机名或IP地址
    cert_file (str): 证书保存路径
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令是否成功执行
        "command": str,       # 执行的命令
        "exit_code": int,     # 退出码
        "stdout": str,        # 标准输出
        "stderr": str         # 标准错误
    }
    """
    cmd = ["gnutls-cli", "--save-cert", cert_file, "--insecure", host]
    result = subprocess.run(cmd, capture_output=True, text=True, input="", timeout=30)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(cmd),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

if __name__ == "__main__":
    mcp.run()